import colorsys
import time

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from matplotlib import pyplot
from torch import nn

from utils.train.imageUtils import letterbox_image, to_categorical, gray2colorfulRgb


class Prediction(object):

    def __init__(self,
                 net,
                 model_path,
                 num_classes,
                 model_image_size=(256, 256, 3),
                 cuda=True,
                 blend=True):
        """
            after improvement, generally need to be modified net, model_path and model_image_size
        """
        self.net = net
        self.model_path = model_path
        self.model_image_size = model_image_size
        self.num_classes = num_classes
        self.cuda = cuda
        self.blend = blend
        self.generate()

    # ---------------------------------------------------#
    #   get all categories
    # ---------------------------------------------------#
    def generate(self):
        state_dict = torch.load(self.model_path)
        self.net.load_state_dict(state_dict)

        if self.cuda:
            self.net = nn.DataParallel(self.net)
            self.net = self.net.cuda()

        print('{} model loaded.'.format(self.model_path))

    # ---------------------------------------------------#
    #   detection picture
    # ---------------------------------------------------#
    def detect_image(self, old_img, images, nw, nh):
        orininal_w, orininal_h = old_img.size

        # ---------------------------------------------------#
        #   Images are transmitted to the network for prediction
        # ---------------------------------------------------#
        with torch.no_grad():
            images = torch.from_numpy(images).type(torch.FloatTensor)
            if self.cuda:
                images = images.cuda()

            pr = self.net(images)
            pr = pr[0]

            # ---------------------------------------------------#
            #   take out the type of each pixel
            # ---------------------------------------------------#
            pr = F.softmax(pr.permute(1, 2, 0), dim=-1).cpu().numpy().argmax(axis=-1)

            # --------------------------------------#
            #   cut off the gray bar
            # --------------------------------------#
            pr = pr[int((self.model_image_size[0] - nh) // 2):int((self.model_image_size[0] - nh) // 2 + nh),
                 int((self.model_image_size[1] - nw) // 2):int((self.model_image_size[1] - nw) // 2 + nw)]

            ############################
            np.savetxt("E:\\360Download\\b.txt", pr)
            temp = Image.fromarray(pr.astype(np.uint8)).convert('L')
            pyplot.imshow(temp)
            pyplot.show()

        image = gray2colorfulRgb(self.num_classes, pr, orininal_w, orininal_h)
        # ------------------------------------------------#
        #   mix the new picture with the original picture
        # ------------------------------------------------#
        if self.blend:
            old_img = old_img.convert('RGB')
            image = Image.blend(old_img, image, 0.7)

        return image

    def detect_image2(self, old_img, images, nw, nh):
        orininal_w, orininal_h = old_img.size

        with torch.no_grad():
            images = torch.from_numpy(images).type(torch.FloatTensor)
            if self.cuda:
                images = images.cuda()

            pr = self.net(images)
            pr = pr[0][0].cpu().numpy()
            image = Image.fromarray(np.uint8(pr)).resize((orininal_w, orininal_h)).convert('L')
            pyplot.imshow(image)
            pyplot.show()

            pr = pr[int((self.model_image_size[0] - nh) // 2):int((self.model_image_size[0] - nh) // 2 + nh),
                 int((self.model_image_size[1] - nw) // 2):int((self.model_image_size[1] - nw) // 2 + nw)]

        image = gray2colorfulRgb(self.num_classes, pr, orininal_w, orininal_h)

        pyplot.imshow(image)
        pyplot.show()

        if self.blend:
            old_img = old_img.convert('RGB')
            image = Image.blend(old_img, image, 0.7)

        return image

    def get_FPS(self, image, test_interval):
        # orininal_h = np.array(image).shape[0]
        # orininal_w = np.array(image).shape[1]

        image, nw, nh = letterbox_image(image, (self.model_image_size[1], self.model_image_size[0]))
        images = [np.array(image) / 255]
        images = np.transpose(images, (0, 3, 1, 2))

        with torch.no_grad():
            images = torch.from_numpy(images).type(torch.FloatTensor)
            if self.cuda:
                images = images.cuda()
            pr = self.net(images)[0]
            pr = F.softmax(pr.permute(1, 2, 0), dim=-1).cpu().numpy().argmax(axis=-1)
            pr = pr[int((self.model_image_size[0] - nh) // 2):int((self.model_image_size[0] - nh) // 2 + nh),
                 int((self.model_image_size[1] - nw) // 2):int((self.model_image_size[1] - nw) // 2 + nw)]

        t1 = time.time()
        for _ in range(test_interval):
            with torch.no_grad():
                pr = self.net(images)[0]
                pr = F.softmax(pr.permute(1, 2, 0), dim=-1).cpu().numpy().argmax(axis=-1)
                pr = pr[int((self.model_image_size[0] - nh) // 2):int((self.model_image_size[0] - nh) // 2 + nh),
                     int((self.model_image_size[1] - nw) // 2):int((self.model_image_size[1] - nw) // 2 + nw)]

        t2 = time.time()
        tact_time = (t2 - t1) / test_interval
        return tact_time
